import cv2
import os
from flask import Flask, render_template, request
from werkzeug.utils import secure_filename
import zipfile


def zip_folder(folder_path, zip_name):
    with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
                
                
# 每间隔20帧取一张图片
def extract_frames(video_path, output_dir, interval=20):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    saved_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        if frame_count % interval == 0:
            output_path = os.path.join(output_dir, f"frame{saved_count}.jpg")
            cv2.imwrite(output_path, frame)
            saved_count += 1

        frame_count += 1

    cap.release()  
    print(f"共提取了{saved_count}张图片")


app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/videos'


@app.route('/')
def index():
    return render_template('index.html')


# conda activate my_project_env
@app.route('/upload', methods=['POST'])
def upload_video():
    video = request.files['video']
    interval = request.form['interval']
    if not os.path.exists(app.config['UPLOAD_FOLDER']):
        os.makedirs(app.config['UPLOAD_FOLDER'])
    if video:
        filename = secure_filename(video.filename)
        video.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))
        video_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        output_dir = "static/output_images/"+video_path
        extract_frames(video_path, output_dir, int(interval))
        zip_folder(output_dir, output_dir+".zip")
        downloadurl = output_dir+".zip"
        return '视频处理成功,每间隔'+interval+'帧取一张图片,下载地址：<a href="'+downloadurl+'">/'+downloadurl+"</a>"
    else:
        return '没有收到视频文件'


if __name__ == '__main__':
    # app.run(debug=True)
    app.run(host='0.0.0.0', port=5000)